# +----------------------------------------------------------------------
# | ChatWork
# +----------------------------------------------------------------------
# | 这不是一个自由软件,您只能在不用于商业目的的前提下对程序代码进行修改和使用。
# | 任何企业和个人不允许对程序代码以任何形式任何目的再发布,商业使用请获取授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
from decimal import Decimal
from typing import Union, List
from fastapi import Query
from pydantic import BaseModel, Field


class KnowDatasetAnnex(BaseModel):
    """ 数据附件 """
    id: str = Field(default="", max_length=64, description="附件ID")
    type: str = Field(..., min_length=1, max_length=32, description="附件类型")
    name: str = Field(..., min_length=1, max_length=200, description="附件名称")
    path: str = Field(..., min_length=1, max_length=250, description="附件路径")
    ext: str = Field(..., min_length=1, max_length=5000, description="附件扩展")
    size: int = Field(..., ge=0, description="附件大小")
    sort: int = Field(..., ge=0, description="附件排序")
    question: str = Field(default="", max_length=5000, description="附件描述")

    class Config:
        json_schema_extra = {
            "example": {
                "id": "32cea0b8-94b4-4f7d-aca1-151e6a930dc6",
                "type": "image",
                "name": "test.png",
                "path": "https://www.xx.com/aa.png",
                "ext": "png",
                "size": 100,
                "sort": 0,
                "question": "附件描述信息"
            }
        }


"""--------------- Separator ---------------"""


class KnowDatasetSearchIn(BaseModel):
    """ 数据搜索参数 """
    page_no: int = Query(default=1, gt=0, description="当前页码")
    page_size: int = Query(default=15, gt=0, le=200, description="每页条数")
    keyword: Union[str, None] = Query(default=None, description="关键词")
    status: Union[int, None] = Query(default=None, ge=0, description="学习状态: [1=排队中,2=学习中,3=学习成功,4=学习失败]")
    archive_id: str = Query(..., min_length=1, max_length=32, description="文档编号")


class KnowDatasetDetailIn(BaseModel):
    """ 数据详情参数 """
    uuids: str = Query(..., description="数据UUID")

    class Config:
        json_schema_extra = {
            "example": {
                "uuids": "0bdbaa93-7905-4ae6-a220-79c8f592cabe"
            }
        }


class KnowDatasetInspectIn(BaseModel):
    """ 数据检测参数 """
    uuids: List[str] = Field(..., description="数据UUID")

    class Config:
        json_schema_extra = {
            "example": {
                "uuids": ["0bdbaa93-7905-4ae6-a220-79c8f592cabe"]
            }
        }


class KnowDatasetInsertIn(BaseModel):
    """ 数据插入参数 """
    fid: str = Field(..., min_length=1, max_length=32, description="文档编号")
    question: str = Field(min_length=1, max_length=5000, description="问题")
    answer: str = Field(max_length=65535, description="内容")
    annex: List[KnowDatasetAnnex] = Field(default=[], description="附件")

    class Config:
        json_schema_extra = {
            "example": {
                "fid": "9bf8d27aed419a085d56ae125c85968f",
                "question": "向量化问题",
                "answer": "附加答案",
                "annex": [
                    {
                        "id": "32cea0b8-94b4-4f7d-aca1-151e6a930dc6",
                        "type": "image",
                        "name": "test.png",
                        "path": "https://www.xx.com/aa.png",
                        "ext": "png",
                        "size": 100,
                        "sort": 0,
                        "question": "附件描述信息"
                    }
                ]
            }
        }


class KnowDatasetUpdateIn(BaseModel):
    """ 数据更新参数 """
    uuid: str = Field(..., min_length=1, max_length=64, description="数据UUID")
    question: str = Field(min_length=1, max_length=5000, description="问题")
    answer: str = Field(max_length=65535, description="内容")
    annex: List[KnowDatasetAnnex] = Field(default=[], description="附件")

    class Config:
        json_schema_extra = {
            "example": {
                "uuid": "0bdbaa93-7905-4ae6-a220-79c8f592cabe",
                "question": "向量化问题",
                "answer": "附加答案",
                "annex": [
                    {
                        "id": "32cea0b8-94b4-4f7d-aca1-151e6a930dc6",
                        "type": "image",
                        "name": "test.png",
                        "path": "https://www.xx.com/aa.png",
                        "ext": "png",
                        "size": 100,
                        "sort": 0,
                        "question": "附件描述信息"
                    }
                ]
            }
        }


class KnowDatasetDeleteIn(BaseModel):
    """ 数据删除参数 """
    kb: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    uuids: List[str] = Field(..., description="数据UUID")

    class Config:
        json_schema_extra = {
            "example": {
                "kb": "wZS3mj23aj6kP95za0",
                "uuids": ["0bdbaa93-7905-4ae6-a220-79c8f592cabe"]
            }
        }


class KnowDatasetRetryIn(BaseModel):
    """ 学习重试参数 """
    uuids: List[str] = Field(..., description="数据UUID")

    class Config:
        json_schema_extra = {
            "example": {
                "uuids": ["0bdbaa93-7905-4ae6-a220-79c8f592cabe"]
            }
        }


class KnowDatasetTestingIn(BaseModel):
    """ 搜索测试参数 """
    kb: str = Field(..., min_length=1, max_length=32, description="知识库编号")
    text: str = Field(..., min_length=1, max_length=9000, description="检索的文本")
    search_mode: str = Field(default="similar", pattern=r"^(similar|full|mix)$", description="检索的模式")
    search_tokens: int = Field(default=5000, ge=0, le=50000, description="检索引用上限")
    search_similar: Decimal = Field(default=0.4, le=1, max_digits=4, decimal_places=3, description="检索相似: [0.0~1.0]")
    ranking_status: int = Field(default=0, ge=0, le=1, description="检索的重排: [0=否, 1=是]")
    ranking_score: Decimal = Field(default=0.5, le=1, max_digits=4, decimal_places=3, description="重排分数: [0.0~1.0]")
    ranking_model: str = Field(default="", max_length=800, description="重排模型")
    optimize_ask: int = Field(default=0, ge=0, le=1, description="问题优化: [0=否, 1=是]")
    optimize_model: str = Field(default="", max_length=800, description="优化模型")

    class Config:
        json_schema_extra = {
            "example": {
                "kb": "wZS3mj23aj6kP95za0",
                "text": "检测的问题",
                "search_mode": "similar",
                "search_tokens": 5000,
                "search_similar": 0.4,
                "ranking_status": 0,
                "optimize_ask": 0,
                "ranking_model": "qianfan:bce-reranker-base",
                "optimize_model": "doubao:doubao-1-5-pro-32k-250115"
            }
        }


"""--------------- Separator ---------------"""


class KnowDatasetListVo(BaseModel):
    """ 数据列表VO """
    id: str = Field(description="UUID")
    question: str = Field(description="问题")
    answer: str = Field(description="答复")
    dimension: int = Field(description="向量维度")
    character: int = Field(description="字符数量")
    use_points: Decimal = Field(description="消耗积分")
    use_tokens: Decimal = Field(description="消耗Tokens")
    status: int = Field(description="学习状态: [1=排队中,2=学习中,3=学习成功,4=学习失败]")
    error: str = Field(description="错误信息")
    create_time: str = Field(description="创建时间")
    update_time: str = Field(description="更新时间")

    class Config:
        json_schema_extra = {
            "example": {
                "id": "fc810deb-05e5-4242-b27e-72d5ab46fa4e",
                "question": "向量化问题",
                "answer": "附加内容",
                "dimension": 0,
                "character": 175,
                "use_points": 0,
                "use_tokens": 0,
                "status": 4,
                "error": "Vector: api_key无效，请管理员检查密钥配置。",
                "create_time": "2025-01-23 14:23:12",
                "update_time": "2025-01-23 14:23:12"
            }
        }


class KnowDatasetDetailVo(BaseModel):
    """ 数据详情VO """
    id: str = Field(description="UUID")
    question: str = Field(description="问题")
    answer: str = Field(description="答复")
    annex: List[KnowDatasetAnnex] = Field(description="附件")

    class Config:
        json_schema_extra = {
            "example": {
                "id": "0bdbaa93-7905-4ae6-a220-79c8f592cabe",
                "question": "",
                "answer": "",
                "annex": [
                    {
                        "id": "32cea0b8-94b4-4f7d-aca1-151e6a930dc6",
                        "type": "image",
                        "name": "test.png",
                        "path": "https://www.xx.com/aa.png",
                        "ext": "png",
                        "size": 100,
                        "sort": 0,
                        "question": "附件描述信息"
                    }
                ]
            }
        }


class KnowDatasetInspectVo(BaseModel):
    """ 数据列表VO """
    id: str = Field(description="UUID")
    use_points: Decimal = Field(description="消耗积分")
    use_tokens: Decimal = Field(description="消耗Tokens")
    status: int = Field(description="训练状态")
    error: str = Field(description="错误信息")
    update_time: str = Field(description="更新时间")

    class Config:
        json_schema_extra = {
            "example": {
                "id": "0bdbaa93-7905-4ae6-a220-79c8f592cabe",
                "use_points": 1,
                "use_tokens": 9,
                "status": 3,
                "error": "",
                "update_time": "2025-01-24 11:34:13"
            }
        }


class KnowDatasetTestingVo(BaseModel):
    """ 搜索测试VO """
    id: str = Field(description="UUID")
    question: str = Field(description="问题")
    answer: str = Field(description="答复")
    character: int = Field(description="字符数量")
    emb_score: float = Field(description="向量分数")
    full_score: float = Field(description="全文分数")
    rerank_score: float = Field(description="重排分数")
    file: dict = Field(description="文件")

    class Config:
        json_schema_extra = {
            "example": {
                "id": "13546169-ec09-45f5-ba34-0f94397e4124",
                "question": "向量化问题",
                "answer": "附加答案",
                "character": 9,
                "emb_score": 0.48962,
                "full_score": 0.48962,
                "rerank_score": 0.48962,
                "file": {
                    "code": "13546169",
                    "name": "aa.pdf",
                    "size": 5682,
                    "path": "https://www.w.com/aa.pdf"
                }
            }
        }
